from lingcod.mpa.models import Mpa as BaseMpa
from lingcod.array.models import MpaArray as BaseMpaArray
from lingcod.raster_stats.models import RasterDataset, zonal_stats, ZonalStatsCache
from lingcod.studyregion.models import StudyRegion
from lingcod.mpa.models import MpaDesignation
from lingcod.unit_converter.models import length_in_display_units, area_in_display_units
from django.contrib.gis.db import models
from django.contrib.auth.models import User, Group
from django.conf import settings
from django.contrib.gis.measure import A
from celery.decorators import task

arczone2 = """+proj=eqc
+lat_0=0.0000000000
+lat_ts=41.1268212700
+lon_0=0.0000000000
+a=6378137
+rf=298.257223563
+no_defs
+towgs84=0.000,0.000,0.000
+to_meter=1.0"""

IMPACT_RASTERS = [
'alt_combined',
'alt_generaluse',
'alt_land10',
'alt_notake',
'original',
'alt_mpabuffer',
'alt_fulluse',
'alt_fishing5',
'alt_climate25'
] # all others assumed to be threats

DesignationRasterLookup = {
  'fulluse': 'alt_fulluse',
  'mpabuffer': 'alt_mpabuffer',
  'generaluse': 'alt_generaluse',
  'notake': 'alt_notake',
  'original': 'original'
}


@task(rate_limit='30/s')
def async_zonal_stats(geom, raster, **kwargs):
    """
    Just wraps the zonal_stats function from raster_stats app 
    Allows it to be run async using celery
    use rate_limit to set how frequently this can run 
    """
    logger = async_zonal_stats.get_logger(**kwargs)
    logger.info("Running zonal stats for %s with a geom %s" % (raster.name,geom.wkt[:60] + "...") )
    
    res = zonal_stats(geom, raster)
    return res

class Zone(BaseMpa):
    pass
    
    @property
    def area_sqmi(self):
        az2_geom = self.geometry_final.transform(arczone2, clone=True)
	return A(sq_m=az2_geom.area).sq_mi
    
    def save(self, *args, **kwargs):
        super(Zone, self).save(*args, **kwargs) # Call the "real" save() method
        self.start_zonal()

    @property
    def impacts(self):
        """
        returns zonal stats for all impact rasters
        ONLY hits the cache and wont
        """
        rasts = RasterDataset.objects.filter(name__in=IMPACT_RASTERS) 
        done = []
        notdone = []
        zoned_rast = 'foo'
        if self.designation and DesignationRasterLookup.has_key(self.designation.acronym):
            zoned_rast = DesignationRasterLookup[self.designation.acronym]
        for rast in rasts:
            rs = zonal_stats(self.geometry_final, rast, cache_only=True)
            if rast.name == zoned_rast:
                rs.zoned = True
            else:
                rs.zoned = False
            done.append(rs)
            if not rs.pixels:
                notdone.append(rast)

        # If its not yet done, fire off another call just in case
        for rast in notdone:
            try:
                z = ZonalStatsCache.objects.get(geom_hash=self.geometry_final.wkt.__hash__(), raster=rast)
                z.delete()
            except:
                pass
            async_zonal_stats.delay(self.geometry_final, rast)

        return done

    @property
    def threats(self):
        """
        returns zonal stats for all nonimpact rasters (threats)
        ONLY hits the cache 
        """
        rasts = RasterDataset.objects.all().exclude(name__in=IMPACT_RASTERS) 
        done = []
        notdone = []
        for rast in rasts:
            rs = zonal_stats(self.geometry_final, rast, cache_only=True)
            done.append(rs)
            if not rs.pixels:
                notdone.append(rast)

        # If its not yet done, fire off another call just in case
        for rast in notdone:
            async_zonal_stats.delay(self.geometry_final, rast)

        return done

    def start_zonal(self):
        """ 
        Async processing for running zonal stats
        returns celery result object list, not the actual zonal results (which are cachec)
        """
        rasts = RasterDataset.objects.all() 
        results = []
        try: 
            x = self.geometry_final.valid
        except:
            # Bad geom
            return None

        for rast in rasts:
            results.append(async_zonal_stats.delay(self.geometry_final, rast))
        return results


class ZoneArray(BaseMpaArray):
    pass

    @property
    def summary_by_designation(self):
        sr_az2_geom = StudyRegion.objects.current().geometry.transform(arczone2, clone=True)
        raw_area = sr_az2_geom.area
        sr_area = area_in_display_units(raw_area)
        by_desig = {}
        desig_list = list(MpaDesignation.objects.all())
        print desig_list
        if self.mpa_set.filter(designation=None):
            desig_list.append(None)
        for desig in desig_list:
            if desig:
                acronym = desig.name
                sort_num = desig.sort
            else:
                acronym = None
                sort_num = 250
            mpas = self.mpa_set.filter(designation=desig)
            count = 0
            summed_area_sq_mi = 0
            for mpa in mpas:
                az2_geom = mpa.geometry_final.transform(arczone2, clone=True)
                raw_area = az2_geom.area
                summed_area_sq_mi += area_in_display_units(raw_area)
                count += 1
                print "   ", mpa
            print summed_area_sq_mi, count
            by_desig[sort_num] = {
                'designation': acronym,
                'count': count,
                'area': summed_area_sq_mi,
                'percent_of_sr': 100 * ( summed_area_sq_mi / sr_area )
            }
        # add in totals for all designations
        count = 0
        summed_area_sq_mi = 0
        for mpa in self.mpa_set:
            az2_geom = mpa.geometry_final.transform(arczone2, clone=True)
            raw_area = az2_geom.area
            summed_area_sq_mi += area_in_display_units(raw_area)
            count += 1
        by_desig[255] = {
            'designation': 'All Zones',
            'count': count,
            'area': summed_area_sq_mi,
            'percent_of_sr': 100 * ( summed_area_sq_mi / sr_area )
        }
        return by_desig

    @property
    def impacts(self):
        """
        returns zonal stats for all impact rasters
        ONLY hits the cache and wont
        """
        rasts = RasterDataset.objects.filter(name__in=IMPACT_RASTERS) 
        raster_results = []
        for rast in rasts:
            tally_sum = 0
            tally_pixels = 0
            for mpa in self.mpa_set:
                rs = zonal_stats(mpa.geometry_final, rast, cache_only=True)
                if not rs.pixels:
                    async_zonal_stats.delay(mpa.geometry_final, rast)
                    #raise Exception("No data for zone %s and raster %s" % (mpa.name, rast.name)) 
                    return None
                else:
                    tally_sum += rs.avg * rs.pixels
                    tally_pixels += rs.pixels

            if rast.name == DesignationRasterLookup['original']:
                raster_results.append( {'rast': 'original', 'avg': tally_sum / tally_pixels} )
            else:
                raster_results.append( {'rast': rast.full_name, 'avg': tally_sum / tally_pixels} )

        tally_sum = 0
        tally_pixels = 0
        for mpa in self.mpa_set:
            if mpa.designation and DesignationRasterLookup.has_key(mpa.designation.acronym):
                rast = RasterDataset.objects.get(name=DesignationRasterLookup[mpa.designation.acronym])
            else:
                #Use current conditions
                rast = RasterDataset.objects.get(name=DesignationRasterLookup['original'])
                
            rs = zonal_stats(mpa.geometry_final, rast, cache_only=True)
            if not rs.pixels:
                async_zonal_stats.delay(mpa.geometry_final, rast)
                return None
            else: 
                tally_sum += rs.avg * rs.pixels
                tally_pixels += rs.pixels

        raster_results.append( {'rast': 'zoned', 'avg': tally_sum / tally_pixels} )

        return raster_results
